#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import PoseStamped
import airsim
import cv2
import numpy as np
from cv_bridge import CvBridge

class AirSimROS2Node(Node):
    def __init__(self):
        super().__init__('airsim_ros2_node')
        self.client = airsim.MultirotorClient()
        self.client.confirmConnection()
        self.client.enableApiControl(True)
        self.client.armDisarm(True)
        self.bridge = CvBridge()
        self.pose_publisher = self.create_publisher(PoseStamped, '/airsim/pose', 10)
        self.image_publisher = self.create_publisher(Image, '/airsim/bottom_camera_image', 10)
        self.image_publisher_ue4 = self.create_publisher(Image, '/airsim/ue4_camera_image', 10)
        self.timer = self.create_timer(0.1, self.timer_callback)
        self.ue4_camera_timer = self.create_timer(0.1, self.ue4_camera_timer_callback)
        self.get_logger().info('AirSimROS2Node started')

    def timer_callback(self):
        # 获取无人机状态
        state = self.client.getMultirotorState()
        pose = PoseStamped()
        pose.header.stamp = self.get_clock().now().to_msg()
        pose.header.frame_id = 'map'
        pose.pose.position.x = state.kinematics_estimated.position.x_val
        pose.pose.position.y = state.kinematics_estimated.position.y_val
        pose.pose.position.z = state.kinematics_estimated.position.z_val
        pose.pose.orientation.x = state.kinematics_estimated.orientation.x_val
        pose.pose.orientation.y = state.kinematics_estimated.orientation.y_val
        pose.pose.orientation.z = state.kinematics_estimated.orientation.z_val
        pose.pose.orientation.w = state.kinematics_estimated.orientation.w_val
        self.pose_publisher.publish(pose)

        # 获取摄像头图像
        responses = self.client.simGetImages([airsim.ImageRequest("BottomCamera", airsim.ImageType.Scene, False, False)])

        if responses:
            img1d = np.frombuffer(responses[0].image_data_uint8, dtype=np.uint8)
            img_rgb = img1d.reshape(responses[0].height, responses[0].width, 3)
            img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
            img_msg = self.bridge.cv2_to_imgmsg(img_bgr, encoding='rgb8')
            self.image_publisher.publish(img_msg)

    
    def ue4_camera_timer_callback(self):
        responses_ue4 = self.client.simGetImages([airsim.ImageRequest("0", airsim.ImageType.Scene, False, False)])
        if responses_ue4:
            img1d_ue4 = np.frombuffer(responses_ue4[0].image_data_uint8, dtype=np.uint8)
            img_rgb_ue4 = img1d_ue4.reshape(responses_ue4[0].height, responses_ue4[0].width, 3)
            img_bgr_ue4 = cv2.cvtColor(img_rgb_ue4, cv2.COLOR_RGB2BGR)
            img_msg_ue4 = self.bridge.cv2_to_imgmsg(img_bgr_ue4, encoding='rgb8')
            self.image_publisher_ue4.publish(img_msg_ue4)

def main(args=None):
    rclpy.init(args=args)
    node = AirSimROS2Node()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
